#include "sendicmp.h"

void sendIcmp::run()
{
    if(!Init())
    {
        return;
    }
    if(!assembleData())
    {
        return;
    }
    u_short seq_no = 0;
    while (!reach_des && max_hop && !is_over)
    {
        max_hop--;
        if(!sendData(seq_no))
        {
            return;
        }
        seq_no++;
        if(!RecvData())
        {
            return;
        }
        ttl++;

    }

    if(!is_over)
    {
        emit send("Trace is comited!");
    }
    else
    {
        emit send("Trace is killed");
    }
    closesocket(raw_sock);
    WSACleanup();
    return;

}

sendIcmp::sendIcmp()
{
    this->timeout = ICMP_MAX_TIMEOUT;
    this->max_hop = ICMP_MAX_HOP;
    this->ip_name = TRACE_EXAMPLE;
    this->padding_letter = PADDING_LETTER;
    this->raw_sock = INVALID_SOCKET;

    this->ttl = 1;

    this->sequence = 0;
    this->time_stamp = 0;
    this->reach_des = false;
    this->is_over = false;
    memset(sender_buffer, 0, sizeof(ICMP_HEADER) + ICMP_DATA_SIZE);
    memset(recv_buffer, 0, ICMP_MAX_LEN);
}

u_short sendIcmp::generateChecksum(u_short *p_buffer, int size)
{
    u_long checksum = 0;
    while(size>1)
    {
        checksum += *p_buffer;
        p_buffer++;
        size -= sizeof(u_short);
    }

    if(size > 0)
    {
        checksum += *(u_char *)p_buffer;
    }

    while(checksum >> 16)
    {
        checksum = (checksum >> 16) + (checksum & 0xFFFF);
    }
    return (u_short)((~checksum)& 0xFFFF);
}

bool sendIcmp::Init()
{
    this->is_over = false;
    this->reach_des  = false;
    this->ttl = 1;
    WSADATA wsa;
    if(WSAStartup(MAKEWORD(2, 2), &wsa) != 0)
    {
        emit send("Fail to initialize the winsock2 dll");
        return false;
    }
    this->raw_sock = WSASocket(AF_INET,
                               SOCK_RAW,
                               IPPROTO_ICMP,
                               NULL,
                               0,
                               WSA_FLAG_OVERLAPPED);

    const char *ip_addr = ip_name.toStdString().c_str();

    ip = inet_addr(ip_addr);

    if(ip == INADDR_NONE)
    {
        hostent *host = gethostbyname(ip_addr);
        if(host)
        {
            ip = (*(in_addr *)host->h_addr).s_addr;
            emit send("Tracing route to [domain] "
                      + QString(inet_ntoa(*(in_addr*)(&ip)))
                      + "\nwith maxmum of " + QString::number(max_hop)
                      + " hops\nwith maxmum timeout " + QString::number(timeout)
                      + " (ms)\n");
        }
        else
        {
            emit send("Fail to resolve host name/ip");
            WSACleanup();
            return false;
        }
    }
    else
    {
        emit send("Tracing route[ip] " + QString(inet_ntoa(*(in_addr*)(&ip)))
                  +" with maxhop " + QString::number(max_hop));

    }
    ZeroMemory(&des_sock_addr, sizeof(sockaddr_in));
    des_sock_addr.sin_family = AF_INET;
    des_sock_addr.sin_addr.s_addr = ip;
    if(raw_sock == INVALID_SOCKET)
    {
        emit send("fail create sock!");
        WSACleanup();
        return false;
    }
    return true;

}

bool sendIcmp::assembleData()
{
    if(setsockopt(raw_sock, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout)) == SOCKET_ERROR)
    {
        emit send(" fail to set recv timeout!");
        closesocket(raw_sock);
        WSACleanup();
        return false;
    }
    icmp = (ICMP_HEADER *)sender_buffer;
    icmp->type = ICMP_REQUEST;
    icmp->code = 0;
    icmp->identification = htons((u_short)(GetCurrentProcessId()));

    memset(sender_buffer + sizeof(ICMP_HEADER), padding_letter, ICMP_DATA_SIZE);
    return true;
}

bool sendIcmp::sendData(u_short seq)
{
    if(setsockopt(raw_sock, IPPROTO_IP, IP_TTL, (char*)&ttl, sizeof(ttl)) == SOCKET_ERROR)
    {
        emit send("Fail to set ttl!");
        closesocket(raw_sock);

        WSACleanup();
        return false;
    }

    ((ICMP_HEADER *)sender_buffer)->sequence = htons(seq);
    ((ICMP_HEADER *)sender_buffer)->checksum = 0;
    ((ICMP_HEADER *)sender_buffer)->checksum = generateChecksum((u_short *)sender_buffer, sizeof(ICMP_HEADER) + ICMP_DATA_SIZE);


    this->sequence = htons(seq);
    this->time_stamp = GetTickCount();

    if(sendto(raw_sock, sender_buffer, sizeof(sender_buffer), 0, (sockaddr *)&des_sock_addr, sizeof(des_sock_addr)) == SOCKET_ERROR)
    {
        if(WSAGetLastError() == WSAEHOSTUNREACH)
        {
            emit send("des host unreachable");
        }
        closesocket(raw_sock);
        WSACleanup();
        return false;
    }
    return true;
}

bool sendIcmp::RecvData()
{
    sockaddr_in from;

    int len = sizeof(from);
    int read_Data_len = 0;

    while(!is_over)
    {
        read_Data_len  = recvfrom(raw_sock, recv_buffer, ICMP_MAX_LEN, 0, (sockaddr *)&from, &len);
        if(read_Data_len != SOCKET_ERROR)
        {
            if(CheckRecvData(recv_buffer, read_Data_len))
            {
                if(recv_sock_addr.sin_addr.s_addr == des_sock_addr.sin_addr.s_addr)
                {
                    reach_des = true;
                    emit send(QString::number(ttl)
                              + " Reach destination "
                              + QString(inet_ntoa(des_sock_addr.sin_addr))
                              + "\n");
                }
                break;
            }
        }
        else if(WSAGetLastError() == WSAETIMEDOUT)
        {
            emit send(QString::number(ttl) + " * Request time out! ");
            break;
        }
        else
        {
            emit send(" Fail to call recvfrom ");
            closesocket(raw_sock);
            WSACleanup();
            return false;
        }
    }
    return true;
}

bool sendIcmp::CheckRecvData(char *recv_buffer, int size)
{
    IP_HEADER *ip_pointer = (IP_HEADER *)(recv_buffer);
    int ip_len = (ip_pointer->version_length & 0x0f) * 4;
    if(size < int(ip_len + sizeof(IP_HEADER)))
    {
        return false;
    }
    icmp = (ICMP_HEADER *)(recv_buffer + ip_len);
    u_short id, seq;
    if(icmp->type == ICMP_RESEPONSE)
    {
        id = ntohs(icmp->identification);
        seq = icmp->sequence;
    }
    else if(icmp->type == ICMP_TTLE)
    {
        char *ip_temp = recv_buffer + ip_len + sizeof(ICMP_HEADER);
        int len = (((IP_HEADER *)ip_temp)->version_length & 0x0f) * 4;
        ICMP_HEADER *icmp_temp = (ICMP_HEADER *)(ip_temp + len);
        id = ntohs(icmp_temp->identification);
        seq = icmp_temp->sequence;
    }
    else
    {
        return false;
    }

    if(id!=(u_short)GetCurrentProcessId() || seq != sequence)
    {
        return false;
    }

    if(icmp->type == ICMP_RESEPONSE || icmp->type == ICMP_TTLE)
    {
        recv_sock_addr.sin_addr.s_addr = ip_pointer->src_addr;
        time_stamp = GetTickCount() - time_stamp;
        if(time_stamp)
        {
            emit send(QString::number(ttl) + " " +
                      QString(inet_ntoa(recv_sock_addr.sin_addr)) + " " + QString::number(time_stamp) + " ms");
        }
        else
        {
            emit send(QString::number(ttl) + " " +
                      QString(inet_ntoa(recv_sock_addr.sin_addr)) + " < 1 ms");
        }
        return true;
    }
    else
    {
        return false;
    }
}

void sendIcmp::setIpName(QString ip_name)
{
    this->ip_name = ip_name;
}

void sendIcmp::setMaxHop(int max_hop)
{
    this->max_hop = max_hop;
}

void sendIcmp::setTimeOut(int timeout)
{
    this->timeout = timeout;
}

void sendIcmp::setKill()
{
    this->is_over = true;
}

void sendIcmp::setPadding(QString letter)
{
    if(letter.size() > 1)
    {
        return;
    }
    else
    {
        this->padding_letter = (letter.toStdString())[0];
        return;
    }
}


